{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "627529e5-9204-4989-855b-fc7193ec17fa",
   "metadata": {},
   "source": [
    "# 🚀 Gradio YOLOv5 Det v0.2 （Jupyter版）\n",
    "## 创建人：曾逸夫\n",
    "## 创建时间：2022-05-08"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1f7190d-4732-44a3-a3b1-cf002dadf007",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fd5defc-74f1-44b1-bf80-baad2a40dded",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.system(\"pip install gradio==2.9.4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cdd5e57-6580-415a-b92d-8bd8e104b495",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import csv\n",
    "import json\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import gradio as gr\n",
    "import torch\n",
    "import yaml\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "\n",
    "from util.fonts_opt import is_fonts\n",
    "from util.pdf_opt import pdf_generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ebffa5-d161-4619-a0ad-70b6631c6670",
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_PATH = sys.path[0]  # 根目录\n",
    "\n",
    "# 本地模型路径\n",
    "local_model_path = f\"{ROOT_PATH}/yolov5\"\n",
    "\n",
    "# Gradio YOLOv5 Det版本\n",
    "GYD_VERSION = \"Gradio YOLOv5 Det v0.3\"\n",
    "\n",
    "# 模型名称临时变量\n",
    "model_name_tmp = \"\"\n",
    "\n",
    "# 设备临时变量\n",
    "device_tmp = \"\"\n",
    "\n",
    "# 文件后缀\n",
    "suffix_list = [\".csv\", \".yaml\"]\n",
    "\n",
    "# 字体大小\n",
    "FONTSIZE = 25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a5c5d0a-7533-44bc-9a44-a01c3a4902cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_args(known=False):\n",
    "    parser = argparse.ArgumentParser(description=\"Gradio YOLOv5 Det v0.2\")\n",
    "    parser.add_argument(\"--model_name\", \"-mn\", default=\"yolov5s\", type=str, help=\"model name\")\n",
    "    parser.add_argument(\n",
    "        \"--model_cfg\",\n",
    "        \"-mc\",\n",
    "        default=\"./model_config/model_name_p5_all.yaml\",\n",
    "        type=str,\n",
    "        help=\"model config\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--cls_name\",\n",
    "        \"-cls\",\n",
    "        default=\"./cls_name/cls_name_zh.yaml\",\n",
    "        type=str,\n",
    "        help=\"cls name\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nms_conf\",\n",
    "        \"-conf\",\n",
    "        default=0.5,\n",
    "        type=float,\n",
    "        help=\"model NMS confidence threshold\",\n",
    "    )\n",
    "    parser.add_argument(\"--nms_iou\", \"-iou\", default=0.45, type=float, help=\"model NMS IoU threshold\")\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--label_dnt_show\",\n",
    "        \"-lds\",\n",
    "        action=\"store_true\",\n",
    "        default=False,\n",
    "        help=\"label show\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--device\",\n",
    "        \"-dev\",\n",
    "        default=\"0\",\n",
    "        type=str,\n",
    "        help=\"cuda or cpu\",\n",
    "    )\n",
    "    parser.add_argument(\"--inference_size\", \"-isz\", default=640, type=int, help=\"model inference size\")\n",
    "\n",
    "    args = parser.parse_known_args()[0] if known else parser.parse_args(args=[])\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4d166ff-689e-4998-9d82-6f9e53d9d60a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# yaml文件解析\n",
    "def yaml_parse(file_path):\n",
    "    return yaml.safe_load(open(file_path, encoding=\"utf-8\").read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd30b37f-df7f-45b0-a76c-0ab838d993d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# yaml csv 文件解析\n",
    "def yaml_csv(file_path, file_tag):\n",
    "    file_suffix = Path(file_path).suffix\n",
    "    if file_suffix == suffix_list[0]:\n",
    "        # 模型名称\n",
    "        file_names = [i[0] for i in list(csv.reader(open(file_path)))]  # csv版\n",
    "    elif file_suffix == suffix_list[1]:\n",
    "        # 模型名称\n",
    "        file_names = yaml_parse(file_path).get(file_tag)  # yaml版\n",
    "    else:\n",
    "        print(f\"{file_path}格式不正确！程序退出！\")\n",
    "        sys.exit()\n",
    "\n",
    "    return file_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66add09b-0a47-45dc-8145-bb1466aca435",
   "metadata": {},
   "outputs": [],
   "source": [
    "#  模型加载\n",
    "def model_loading(model_name, device):\n",
    "\n",
    "    # 加载本地模型\n",
    "    model = torch.hub.load(\n",
    "        local_model_path,\n",
    "        \"custom\",\n",
    "        path=f\"{local_model_path}/{model_name}\",\n",
    "        source=\"local\",\n",
    "        device=device,\n",
    "        _verbose=False,\n",
    "    )\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6523632-927f-43a1-b656-2094662a4f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 检测信息\n",
    "def export_json(results, model, img_size):\n",
    "\n",
    "    return [\n",
    "        [\n",
    "            {\n",
    "                \"id\": i,\n",
    "                \"class\": int(result[i][5]),\n",
    "                # \"class_name\": model.model.names[int(result[i][5])],\n",
    "                \"class_name\": model_cls_name_cp[int(result[i][5])],\n",
    "                \"normalized_box\": {\n",
    "                    \"x0\": round(result[i][:4].tolist()[0], 6),\n",
    "                    \"y0\": round(result[i][:4].tolist()[1], 6),\n",
    "                    \"x1\": round(result[i][:4].tolist()[2], 6),\n",
    "                    \"y1\": round(result[i][:4].tolist()[3], 6),\n",
    "                },\n",
    "                \"confidence\": round(float(result[i][4]), 2),\n",
    "                \"fps\": round(1000 / float(results.t[1]), 2),\n",
    "                \"width\": img_size[0],\n",
    "                \"height\": img_size[1],\n",
    "            }\n",
    "            for i in range(len(result))\n",
    "        ]\n",
    "        for result in results.xyxyn\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "375f9250-9952-484a-95b6-dba8e4eefb60",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 帧转换\n",
    "def pil_draw(img, countdown_msg, textFont, xyxy, font_size, label_opt):\n",
    "\n",
    "    img_pil = ImageDraw.Draw(img)\n",
    "\n",
    "    img_pil.rectangle(xyxy, fill=None, outline=\"green\")  # 边界框\n",
    "\n",
    "    if label_opt:\n",
    "        text_w, text_h = textFont.getsize(countdown_msg)  # 标签尺寸\n",
    "        img_pil.rectangle(\n",
    "            (xyxy[0], xyxy[1], xyxy[0] + text_w, xyxy[1] + text_h),\n",
    "            fill=\"green\",\n",
    "            outline=\"green\",\n",
    "        )  # 标签背景\n",
    "        img_pil.multiline_text(\n",
    "            (xyxy[0], xyxy[1]),\n",
    "            countdown_msg,\n",
    "            fill=(205, 250, 255),\n",
    "            font=textFont,\n",
    "            align=\"center\",\n",
    "        )\n",
    "\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9078a7-c36f-4b03-b1ff-596fdb457681",
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOLOv5图片检测函数\n",
    "def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, model_cls, opt):\n",
    "\n",
    "    global model, model_name_tmp, device_tmp\n",
    "\n",
    "    if model_name_tmp != model_name:\n",
    "        # 模型判断，避免反复加载\n",
    "        model_name_tmp = model_name\n",
    "        model = model_loading(model_name_tmp, device)\n",
    "    elif device_tmp != device:\n",
    "        device_tmp = device\n",
    "        model = model_loading(model_name_tmp, device)\n",
    "\n",
    "    # -----------模型调参-----------\n",
    "    model.conf = conf  # NMS 置信度阈值\n",
    "    model.iou = iou  # NMS IOU阈值\n",
    "    model.max_det = 1000  # 最大检测框数\n",
    "    model.classes = model_cls  # 模型类别\n",
    "\n",
    "    results = model(img, size=inference_size)  # 检测\n",
    "\n",
    "    img_size = img.size  # 帧尺寸\n",
    "\n",
    "    # ----------------加载字体----------------\n",
    "    yaml_index = cls_name.index(\".yaml\")\n",
    "    cls_name_lang = cls_name[yaml_index - 2:yaml_index]\n",
    "\n",
    "    if cls_name_lang == \"zh\":\n",
    "        # 中文\n",
    "        textFont = ImageFont.truetype(str(f\"{ROOT_PATH}/fonts/SimSun.ttf\"), size=FONTSIZE)\n",
    "    elif cls_name_lang in [\"en\", \"ru\", \"es\", \"ar\"]:\n",
    "        # 英文、俄语、西班牙语、阿拉伯语\n",
    "        textFont = ImageFont.truetype(str(f\"{ROOT_PATH}/fonts/TimesNewRoman.ttf\"), size=FONTSIZE)\n",
    "    elif cls_name_lang == \"ko\":\n",
    "        # 韩语\n",
    "        textFont = ImageFont.truetype(str(f\"{ROOT_PATH}/fonts/malgun.ttf\"), size=FONTSIZE)\n",
    "\n",
    "    det_img = img.copy()\n",
    "\n",
    "    for result in results.xyxyn:\n",
    "        for i in range(len(result)):\n",
    "            id = int(i)  # 实例ID\n",
    "            obj_cls_index = int(result[i][5])  # 类别索引\n",
    "            obj_cls = model_cls_name_cp[obj_cls_index]  # 类别\n",
    "\n",
    "            # ------------边框坐标------------\n",
    "            x0 = float(result[i][:4].tolist()[0])\n",
    "            y0 = float(result[i][:4].tolist()[1])\n",
    "            x1 = float(result[i][:4].tolist()[2])\n",
    "            y1 = float(result[i][:4].tolist()[3])\n",
    "\n",
    "            # ------------边框实际坐标------------\n",
    "            x0 = int(img_size[0] * x0)\n",
    "            y0 = int(img_size[1] * y0)\n",
    "            x1 = int(img_size[0] * x1)\n",
    "            y1 = int(img_size[1] * y1)\n",
    "\n",
    "            conf = float(result[i][4])  # 置信度\n",
    "            # fps = f\"{(1000 / float(results.t[1])):.2f}\"  # FPS\n",
    "\n",
    "            det_img = pil_draw(\n",
    "                img,\n",
    "                f\"{id}-{obj_cls}:{conf:.2f}\",\n",
    "                textFont,\n",
    "                [x0, y0, x1, y1],\n",
    "                FONTSIZE,\n",
    "                label_opt,\n",
    "            )\n",
    "\n",
    "    det_json = export_json(results, model, img.size)[0]  # 检测信息\n",
    "\n",
    "    # JSON格式化\n",
    "    det_json_format = json.dumps(det_json, sort_keys=True, indent=4, separators=(\",\", \":\"), ensure_ascii=False)\n",
    "\n",
    "    # -------pdf-------\n",
    "    report = \"./Det_Report.pdf\"\n",
    "    if \"pdf\" in opt:\n",
    "        pdf_generate(f\"{det_json_format}\", report, GYD_VERSION)\n",
    "    else:\n",
    "        report = None\n",
    "\n",
    "    if \"json\" not in opt:\n",
    "        det_json = None\n",
    "\n",
    "    return det_img, det_json, report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "502b4d1c-d628-4d57-bfa7-1ff94ba835a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(args):\n",
    "    gr.close_all()\n",
    "\n",
    "    global model, model_cls_name_cp, cls_name\n",
    "\n",
    "    slider_step = 0.05  # 滑动步长\n",
    "\n",
    "    nms_conf = args.nms_conf\n",
    "    nms_iou = args.nms_iou\n",
    "    label_opt = args.label_dnt_show\n",
    "    model_name = args.model_name\n",
    "    model_cfg = args.model_cfg\n",
    "    cls_name = args.cls_name\n",
    "    device = args.device\n",
    "    inference_size = args.inference_size\n",
    "\n",
    "    is_fonts(f\"{ROOT_PATH}/fonts\")  # 检查字体文件\n",
    "\n",
    "    # 模型加载\n",
    "    model = model_loading(model_name, device)\n",
    "\n",
    "    model_names = yaml_csv(model_cfg, \"model_names\")\n",
    "    model_cls_name = yaml_csv(cls_name, \"model_cls_name\")\n",
    "\n",
    "    model_cls_name_cp = model_cls_name.copy()  # 类别名称\n",
    "\n",
    "    # -------------------输入组件-------------------\n",
    "    inputs_img = gr.inputs.Image(type=\"pil\", label=\"原始图片\")\n",
    "    inputs_device = gr.inputs.Dropdown(choices=[\"0\", \"cpu\"], default=device, type=\"value\", label=\"设备\")\n",
    "    inputs_model = gr.inputs.Dropdown(choices=model_names, default=model_name, type=\"value\", label=\"模型\")\n",
    "    inputs_size = gr.inputs.Radio(choices=[320, 640], default=inference_size, label=\"推理尺寸\")\n",
    "    input_conf = gr.inputs.Slider(0, 1, step=slider_step, default=nms_conf, label=\"置信度阈值\")\n",
    "    inputs_iou = gr.inputs.Slider(0, 1, step=slider_step, default=nms_iou, label=\"IoU 阈值\")\n",
    "    inputs_label = gr.inputs.Checkbox(default=(not label_opt), label=\"标签显示\")\n",
    "    inputs_clsName = gr.inputs.CheckboxGroup(choices=model_cls_name, default=model_cls_name, type=\"index\", label=\"类别\")\n",
    "    inputs_opt = gr.inputs.CheckboxGroup(choices=[\"pdf\", \"json\"], default=[\"pdf\"], type=\"value\", label=\"操作\")\n",
    "\n",
    "    # 输入参数\n",
    "    inputs = [\n",
    "        inputs_img,  # 输入图片\n",
    "        inputs_device,  # 设备\n",
    "        inputs_model,  # 模型\n",
    "        inputs_size,  # 推理尺寸\n",
    "        input_conf,  # 置信度阈值\n",
    "        inputs_iou,  # IoU阈值\n",
    "        inputs_label,  # 标签显示\n",
    "        inputs_clsName,  # 类别\n",
    "        inputs_opt,  # 检测操作\n",
    "    ]\n",
    "\n",
    "    # 输出参数\n",
    "    outputs_img = gr.outputs.Image(type=\"pil\", label=\"检测图片\")\n",
    "    outputs02_json = gr.outputs.JSON(label=\"检测信息\")\n",
    "    outputs03_pdf = gr.outputs.File(label=\"下载检测报告\")\n",
    "\n",
    "    outputs = [outputs_img, outputs02_json, outputs03_pdf]\n",
    "\n",
    "    # 标题\n",
    "    title = \"基于Gradio的YOLOv5通用目标检测系统v0.2\"\n",
    "\n",
    "    # 描述\n",
    "    description = \"<div align='center'>可自定义目标检测模型、安装简单、使用方便</div>\"\n",
    "\n",
    "    # 示例图片\n",
    "    examples = [\n",
    "        [\n",
    "            \"./img_example/bus.jpg\",\n",
    "            \"cpu\",\n",
    "            \"yolov5s\",\n",
    "            640,\n",
    "            0.6,\n",
    "            0.5,\n",
    "            True,\n",
    "            [\"人\", \"公交车\"],\n",
    "            [\"pdf\"],],\n",
    "        [\n",
    "            \"./img_example/Millenial-at-work.jpg\",\n",
    "            \"0\",\n",
    "            \"yolov5l\",\n",
    "            320,\n",
    "            0.5,\n",
    "            0.45,\n",
    "            True,\n",
    "            [\"人\", \"椅子\", \"杯子\", \"笔记本电脑\"],\n",
    "            [\"json\"],],\n",
    "        [\n",
    "            \"./img_example/zidane.jpg\",\n",
    "            \"0\",\n",
    "            \"yolov5m\",\n",
    "            640,\n",
    "            0.25,\n",
    "            0.5,\n",
    "            False,\n",
    "            [\"人\", \"领带\"],\n",
    "            [\"pdf\", \"json\"],],]\n",
    "\n",
    "    # 接口\n",
    "    gr.Interface(\n",
    "        fn=yolo_det,\n",
    "        inputs=inputs,\n",
    "        outputs=outputs,\n",
    "        title=title,\n",
    "        description=description,\n",
    "        examples=examples,\n",
    "        theme=\"seafoam\",\n",
    "        # live=True, # 实时变更输出\n",
    "        flagging_dir=\"run\",  # 输出目录\n",
    "        # flagging_options=[\"good\", \"generally\", \"bad\"],\n",
    "        # allow_flagging=\"auto\",\n",
    "        # ).launch(inbrowser=True, auth=['admin', 'admin'])\n",
    "    ).launch(\n",
    "        inbrowser=True,  # 自动打开默认浏览器\n",
    "        show_tips=True,  # 自动显示gradio最新功能\n",
    "        # favicon_path=\"./icon/logo.ico\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b976790e-77ef-4332-b2a4-1d2d5bcf2cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    args = parse_args()\n",
    "    main(args)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f9d149eed12fac937545f4635efc9e9926f24a5acb73f9a6fbfbd1c202313110"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
